#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
传感器数据提取器
专门用于从S1手环原始数据中提取EDA和PPG数据
"""

import pandas as pd
import numpy as np
from pathlib import Path
from typing import Union, Optional, Tuple, List
import logging
import re
import json
from datetime import datetime

# 设置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class SensorExtractor:
    """传感器数据提取器"""
    
    def __init__(self):
        """初始化提取器"""
        # 传感器采样率映射
        self.sensor_sampling_rates = {
            'EDA': 4,    # EDA采样率4Hz
            'PPG': 25,   # PPG采样率25Hz (对应HRP_25)
        }
        
        # 传感器在数据中的位置映射
        self.sensor_positions = {
            'EDA': 2,    # EDA_4在第3个位置（索引2）
            'PPG': 3,    # HRP_25在第4个位置（索引3）
        }
    
    def parse_raw_file(self, file_path: Union[str, Path]) -> Tuple[dict, int, List[str]]:
        """
        解析原始数据文件，提取传感器信息和数据行
        
        Args:
            file_path: 原始数据文件路径
            
        Returns:
            tuple: (传感器信息字典, 开始时间戳, 数据行列表)
        """
        file_path = Path(file_path)
        
        if not file_path.exists():
            raise FileNotFoundError(f"文件不存在: {file_path}")
        
        logger.info(f"正在解析原始数据文件: {file_path}")
        
        with open(file_path, 'r', encoding='utf-8') as f:
            lines = f.readlines()
        
        if len(lines) < 3:
            raise ValueError("文件格式不正确，至少需要3行数据")
        
        # 解析第一行传感器信息
        sensor_line = lines[0].strip()
        sensors = {}
        sensor_parts = sensor_line.split(', ')
        
        for part in sensor_parts:
            if '_' in part:
                sensor_name, rate = part.split('_')
                sensors[sensor_name] = int(rate)
        
        logger.info(f"传感器信息: {sensors}")
        
        # 解析第二行开始时间
        start_time_line = lines[1].strip()
        start_time_match = re.search(r'start time:(\d+)', start_time_line)
        if not start_time_match:
            raise ValueError("无法解析开始时间")
        
        start_time = int(start_time_match.group(1))
        logger.info(f"开始时间: {start_time}")
        
        # 获取数据行（从第3行开始，排除最后一行的结束时间）
        data_lines = []
        for line in lines[2:]:
            line = line.strip()
            if line and not line.startswith('stop time:'):
                data_lines.append(line)
        
        logger.info(f"数据行数: {len(data_lines)}")
        
        return sensors, start_time, data_lines
    
    def extract_sensor_data(self, data_lines: List[str], sensor_name: str, 
                          sensor_position: int, sampling_rate: int, 
                          start_time: int) -> List[Tuple[int, float]]:
        """
        从数据行中提取指定传感器的数据
        
        Args:
            data_lines: 数据行列表
            sensor_name: 传感器名称
            sensor_position: 传感器在数据中的位置
            sampling_rate: 采样率
            start_time: 开始时间戳
            
        Returns:
            list: [(时间戳(ms), 传感器值), ...]
        """
        logger.info(f"正在提取{sensor_name}数据，采样率: {sampling_rate}Hz")
        
        extracted_data = []
        
        for line_idx, line in enumerate(data_lines):
            # 解析时间戳
            parts = line.split(' ', 1)
            if len(parts) < 2:
                continue
                
            timestamp = int(parts[0])
            data_part = parts[1]
            
            # 按&分割传感器数据
            sensor_blocks = data_part.split('&')
            
            if len(sensor_blocks) <= sensor_position:
                logger.warning(f"第{line_idx+1}行数据格式不正确，跳过")
                continue
            
            # 获取指定传感器的数据块
            sensor_block = sensor_blocks[sensor_position].strip()
            
            if not sensor_block:
                continue
            
            # 按,分割一秒内的多个采样点
            samples = sensor_block.split(',')
            
            # 直接使用源文件的时间戳，不进行插值处理
            for sample_idx, sample in enumerate(samples):
                if not sample.strip():
                    continue
                
                try:
                    # 对于EDA，取第一个值；对于PPG，也取第一个值
                    sample_values = sample.split('@')
                    if sample_values:
                        sensor_value = float(sample_values[0])
                        
                        # 直接使用源文件的时间戳（转换为毫秒）
                        sample_timestamp_ms = timestamp * 1000
                        
                        extracted_data.append((sample_timestamp_ms, sensor_value))
                        
                except (ValueError, IndexError) as e:
                    logger.warning(f"解析采样点数据失败: {sample}, 错误: {e}")
                    continue
        
        logger.info(f"{sensor_name}数据提取完成，共{len(extracted_data)}个数据点")
        return extracted_data
    
    def extract_eda_data(self, file_path: Union[str, Path], 
                        include_timestamp: bool = True) -> pd.DataFrame:
        """
        提取EDA数据
        
        Args:
            file_path: 原始数据文件路径
            include_timestamp: 是否包含时间戳列
            
        Returns:
            pd.DataFrame: EDA数据
        """
        sensors, start_time, data_lines = self.parse_raw_file(file_path)
        
        # 检查是否有EDA传感器
        if 'EDA' not in sensors:
            raise ValueError("原始数据中未找到EDA传感器")
        
        # 提取EDA数据
        eda_data = self.extract_sensor_data(
            data_lines, 'EDA', self.sensor_positions['EDA'], 
            self.sensor_sampling_rates['EDA'], start_time
        )
        
        # 创建DataFrame
        if include_timestamp:
            df = pd.DataFrame(eda_data, columns=['timestamp', 'EDA'])
        else:
            df = pd.DataFrame([value for _, value in eda_data], columns=['EDA'])
        
        return df
    
    def extract_ppg_data(self, file_path: Union[str, Path], 
                        include_timestamp: bool = True) -> pd.DataFrame:
        """
        提取PPG数据
        
        Args:
            file_path: 原始数据文件路径
            include_timestamp: 是否包含时间戳列
            
        Returns:
            pd.DataFrame: PPG数据
        """
        sensors, start_time, data_lines = self.parse_raw_file(file_path)
        
        # 检查是否有PPG传感器（对应HRP）
        if 'HRP' not in sensors:
            raise ValueError("原始数据中未找到PPG传感器(HRP)")
        
        # 提取PPG数据
        ppg_data = self.extract_sensor_data(
            data_lines, 'PPG', self.sensor_positions['PPG'], 
            self.sensor_sampling_rates['PPG'], start_time
        )
        
        # 创建DataFrame
        if include_timestamp:
            df = pd.DataFrame(ppg_data, columns=['timestamp', 'PPG'])
        else:
            df = pd.DataFrame([value for _, value in ppg_data], columns=['PPG'])
        
        return df
    
    def save_eda_csv(self, file_path: Union[str, Path], output_dir: Union[str, Path] = None,
                    include_timestamp: bool = False) -> str:
        """
        提取并保存EDA数据为CSV文件
        
        Args:
            file_path: 原始数据文件路径
            output_dir: 输出目录，默认为当前目录下的output文件夹
            include_timestamp: 是否包含时间戳列，默认不包含
            
        Returns:
            str: 输出文件路径
        """
        file_path = Path(file_path)
        
        if output_dir is None:
            output_dir = Path('output')
        else:
            output_dir = Path(output_dir)
        
        output_dir.mkdir(parents=True, exist_ok=True)
        
        # 生成输出文件名
        base_name = file_path.stem
        output_file = output_dir / f"{base_name}_EDA.csv"
        
        # 提取EDA数据
        eda_df = self.extract_eda_data(file_path, include_timestamp)
        
        # 按照新格式保存CSV文件
        with open(output_file, 'w', encoding='utf-8') as f:

            # 写入数据
            if include_timestamp:
                # 含时间戳：两列
                f.write("timestamp,EDA\n")
                for _, row in eda_df.iterrows():
                    f.write(f"{int(row['timestamp'])},{row['EDA']}\n")
            else:
                # 不含时间戳：一列
                f.write("EDA\n")
                for _, row in eda_df.iterrows():
                    f.write(f"{row['EDA']}\n")
        
        logger.info(f"EDA数据已保存到: {output_file}")
        
        return str(output_file)
    
    def save_ppg_csv(self, file_path: Union[str, Path], output_dir: Union[str, Path] = None,
                    include_timestamp: bool = False) -> str:
        """
        提取并保存PPG数据为CSV文件
        
        Args:
            file_path: 原始数据文件路径
            output_dir: 输出目录，默认为当前目录下的output文件夹
            include_timestamp: 是否包含时间戳列，默认不包含
            
        Returns:
            str: 输出文件路径
        """
        file_path = Path(file_path)
        
        if output_dir is None:
            output_dir = Path('output')
        else:
            output_dir = Path(output_dir)
        
        output_dir.mkdir(parents=True, exist_ok=True)
        
        # 生成输出文件名
        base_name = file_path.stem
        output_file = output_dir / f"{base_name}_PPG.csv"
        
        # 提取PPG数据
        ppg_df = self.extract_ppg_data(file_path, include_timestamp)
        
        # 按照新格式保存CSV文件
        with open(output_file, 'w', encoding='utf-8') as f:
            
            # 写入数据
            if include_timestamp:
                # 含时间戳：两列
                f.write("timestamp,PPG\n")
                for _, row in ppg_df.iterrows():
                    f.write(f"{int(row['timestamp'])},{row['PPG']}\n")
            else:
                # 不含时间戳：一列
                f.write("PPG\n")
                for _, row in ppg_df.iterrows():
                    f.write(f"{row['PPG']}\n")
        
        logger.info(f"PPG数据已保存到: {output_file}")
        return str(output_file)
    
    def save_eda_json(self, file_path: Union[str, Path], output_dir: Union[str, Path] = None,
                     include_timestamp: bool = False) -> str:
        """
        提取EDA数据并保存为JSON格式
        
        Args:
            file_path: 原始数据文件路径
            output_dir: 输出目录，默认为当前目录
            include_timestamp: 是否包含时间戳
            
        Returns:
            str: 输出文件路径
        """
        # 提取EDA数据
        eda_df = self.extract_eda_data(file_path, include_timestamp=True)
        
        if eda_df.empty:
            raise ValueError("未找到EDA数据")
        
        # 准备输出路径
        file_path = Path(file_path)
        if output_dir is None:
            output_dir = file_path.parent
        else:
            output_dir = Path(output_dir)
            output_dir.mkdir(parents=True, exist_ok=True)
        
        # 生成输出文件名
        base_name = file_path.stem
        output_file = output_dir / f"{base_name}_EDA.json"
        
        # 计算时间信息
        start_timestamp = eda_df['timestamp'].iloc[0] / 1000  # 转换为秒
        end_timestamp = eda_df['timestamp'].iloc[-1] / 1000
        duration = end_timestamp - start_timestamp
        sampling_rate = self.sensor_sampling_rates['EDA']
        
        # 转换时间戳为ISO格式
        start_time = datetime.fromtimestamp(start_timestamp).isoformat() + "+08:00"
        end_time = datetime.fromtimestamp(end_timestamp).isoformat() + "+08:00"
        
        # 准备JSON数据
        json_data = {
            "start_time": start_time,
            "end_time": end_time,
            "duration": duration,
            "sampling_rate": sampling_rate
        }
        
        if include_timestamp:
            # 包含时间戳的格式
            json_data["data"] = [
                {
                    "timestamp": (row['timestamp'] - eda_df['timestamp'].iloc[0]) / 1000,  # 相对时间
                    "value": row['EDA']
                }
                for _, row in eda_df.iterrows()
            ]
        else:
            # 不包含时间戳的格式
            json_data["data"] = eda_df['EDA'].tolist()
        
        # 保存JSON文件
        with open(output_file, 'w', encoding='utf-8') as f:
            json.dump(json_data, f, indent=4, ensure_ascii=False)
        
        logger.info(f"EDA数据已保存到: {output_file}")
        return str(output_file)
    
    def save_ppg_json(self, file_path: Union[str, Path], output_dir: Union[str, Path] = None,
                     include_timestamp: bool = False) -> str:
        """
        提取PPG数据并保存为JSON格式
        
        Args:
            file_path: 原始数据文件路径
            output_dir: 输出目录，默认为当前目录
            include_timestamp: 是否包含时间戳
            
        Returns:
            str: 输出文件路径
        """
        # 提取PPG数据
        ppg_df = self.extract_ppg_data(file_path, include_timestamp=True)
        
        if ppg_df.empty:
            raise ValueError("未找到PPG数据")
        
        # 准备输出路径
        file_path = Path(file_path)
        if output_dir is None:
            output_dir = file_path.parent
        else:
            output_dir = Path(output_dir)
            output_dir.mkdir(parents=True, exist_ok=True)
        
        # 生成输出文件名
        base_name = file_path.stem
        output_file = output_dir / f"{base_name}_PPG.json"
        
        # 计算时间信息
        start_timestamp = ppg_df['timestamp'].iloc[0] / 1000  # 转换为秒
        end_timestamp = ppg_df['timestamp'].iloc[-1] / 1000
        duration = end_timestamp - start_timestamp
        sampling_rate = self.sensor_sampling_rates['PPG']
        
        # 转换时间戳为ISO格式
        start_time = datetime.fromtimestamp(start_timestamp).isoformat() + "+08:00"
        end_time = datetime.fromtimestamp(end_timestamp).isoformat() + "+08:00"
        
        # 准备JSON数据
        json_data = {
            "start_time": start_time,
            "end_time": end_time,
            "duration": duration,
            "sampling_rate": sampling_rate
        }
        
        if include_timestamp:
            # 包含时间戳的格式
            json_data["data"] = [
                {
                    "timestamp": (row['timestamp'] - ppg_df['timestamp'].iloc[0]) / 1000,  # 相对时间
                    "value": row['PPG']
                }
                for _, row in ppg_df.iterrows()
            ]
        else:
            # 不包含时间戳的格式
            json_data["data"] = ppg_df['PPG'].tolist()
        
        # 保存JSON文件
        with open(output_file, 'w', encoding='utf-8') as f:
            json.dump(json_data, f, indent=4, ensure_ascii=False)
        
        logger.info(f"PPG数据已保存到: {output_file}")
        return str(output_file)


def main():
    """主函数示例"""
    extractor = SensorExtractor()
    
    # 示例用法
    input_file = ".test-data/冰川1后d9221935c5f7_1757232978.txt"
    
    try:
        # 提取并保存EDA数据
        eda_file = extractor.save_eda_csv(input_file, include_timestamp=True)
        print(f"EDA数据已保存到: {eda_file}")
        
        # 提取并保存PPG数据
        ppg_file = extractor.save_ppg_csv(input_file, include_timestamp=True)
        print(f"PPG数据已保存到: {ppg_file}")
        
    except Exception as e:
        logger.error(f"处理过程中发生错误: {e}")


if __name__ == "__main__":
    main()